[OpenCV] Pytorchの手書き数字(MNIST)分類モデルをOpenCVから利用してみました
1 はじめに
CX事業本部の平内(SIN)です。
Pytorch入門ということで、MNIST(手書き数字のデータセット)から作成したモデルを使用して、OpenCVでWebカメラの動画を推論にかけてみました。
使用したモデルのコードは、Githubで公開されている、Pytorchの公式サンプルコードです。
https://github.com/pytorch/examples/blob/master/mnist/main.py
最初に動作しているようすです。 推論の対象となっているのは、画面の中央だけで、別ウインドウに表示されている部分です。
2 モデル
公開されているサンプルコードは、以下のようになっています。
(1) データセット
データセットは、torchvisionによって、MNISTが利用されています。 取得時に、transformsによる変換を行って、訓練用とテスト用のデータローダーが準備されています。
ちなみに、コマンドラインからデフォルト値で使用した場合、学習用のミニバッチは、64となります。
from torchvision import datasets, transforms transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) dataset1 = datasets.MNIST('../data', train=True, download=True, transform=transform) dataset2 = datasets.MNIST('../data', train=False, transform=transform) train_loader = torch.utils.data.DataLoader(dataset1, **kwargs) test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)
(2) モデル
入力は、畳み込み層✕2、プーリング層✕1で処理された後、全結合層(9216 -> 128 -> 10)を通過します。また、途中で、ドロップアウトも挟まれています。
class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) self.conv2 = nn.Conv2d(32, 64, 3, 1) self.dropout1 = nn.Dropout2d(0.25) self.dropout2 = nn.Dropout2d(0.5) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = F.relu(x) x = self.conv2(x) x = F.relu(x) x = F.max_pool2d(x, 2) x = self.dropout1(x) x = torch.flatten(x, 1) x = self.fc1(x) x = F.relu(x) x = self.dropout2(x) x = self.fc2(x) output = F.log_softmax(x, dim=1) return output
(3) 学習
オプティマイザは、optim.Adadelta()で、学習レート(lr)は、デフォルト値で1.0となっています。 また、損失関数 F.nll_loss()です。
optimizer = optim.Adadelta(model.parameters(), lr=args.lr) F.nll_loss(output, target)
下記は、1エポックの中で呼ばれている訓練用のコードです。 ミニバッチごとにデータローダから訓練用データを取り出し、学習が進められています。
def train(args, model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % args.log_interval == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) if args.dry_run: break
コマンドラインから、下記のように利用することで、エポック14回(デフォルト値)、学習が進みます。 手元のMBP(2.3GHz Core i5)では、1エポックに、約2分強で、14エポック終了まで約30分程度かかりました。
テストデータの検証結果を見た感じ、epochは、7回ぐらいで良いのかも知れません。
% python3 index.py --log-interval=100 --save-model Train Epoch: 1 [0/60000 (0%)] Loss: 2.293032 Train Epoch: 1 [6400/60000 (11%)] Loss: 0.273636 Train Epoch: 1 [12800/60000 (21%)] Loss: 0.294438 Train Epoch: 1 [19200/60000 (32%)] Loss: 0.360903 Train Epoch: 1 [25600/60000 (43%)] Loss: 0.097714 Train Epoch: 1 [32000/60000 (53%)] Loss: 0.270267 Train Epoch: 1 [38400/60000 (64%)] Loss: 0.145856 Train Epoch: 1 [44800/60000 (75%)] Loss: 0.173067 Train Epoch: 1 [51200/60000 (85%)] Loss: 0.289421 Train Epoch: 1 [57600/60000 (96%)] Loss: 0.154544 Test set: Average loss: 0.0497, Accuracy: 9835/10000 (98%) ・・・略・・・ Train Epoch: 2 Test set: Average loss: 0.0409, Accuracy: 9871/10000 (99%) Train Epoch: 3 Test set: Average loss: 0.0354, Accuracy: 9876/10000 (99%) Train Epoch: 4 Test set: Average loss: 0.0342, Accuracy: 9886/10000 (99%) Train Epoch: 5 Test set: Average loss: 0.0325, Accuracy: 9893/10000 (99%) Train Epoch: 6 Test set: Average loss: 0.0317, Accuracy: 9896/10000 (99%) Train Epoch: 7 Test set: Average loss: 0.0303, Accuracy: 9900/10000 (99%) Train Epoch: 8 Test set: Average loss: 0.0308, Accuracy: 9901/10000 (99%) Train Epoch: 9 Test set: Average loss: 0.0298, Accuracy: 9905/10000 (99%) Train Epoch: 10 Test set: Average loss: 0.0294, Accuracy: 9910/10000 (99%) Train Epoch: 11 Test set: Average loss: 0.0292, Accuracy: 9910/10000 (99%) Train Epoch: 12 Test set: Average loss: 0.0287, Accuracy: 9911/10000 (99%) Train Epoch: 13 Test set: Average loss: 0.0287, Accuracy: 9910/10000 (99%) Train Epoch: 14 Test set: Average loss: 0.0286, Accuracy: 9908/10000 (99%)
--save-modelのスイッチを付けておくと、学習終了時にモデル(mnist_cnn.pt)が保存されます。
% ls -la *.pt -rw-r--r-- 1 hirauchi.shinichi staff 4800893 7 19 02:58 mnist_cnn.pt
3 入力形式
モデルの入力形式を確認するために、データローダからデータを1つ取り出して確認しました。
import torch from torchvision import datasets, transforms import numpy as np import matplotlib.pyplot as plt transform=transforms.Compose([ transforms.ToTensor() ]) test_data = datasets.MNIST('../data', train=False, download=True, transform=transform) test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=False) tmp = test_loader.__iter__() data, labels = tmp.next() # size print(data.size()) # torch.Size([1, 1, 28, 28]) # imshow img = data.numpy().reshape((28, 28)) plt.imshow(img, cmap='gray') # dump print(data)
(1) size()
学習時にモデルに送られるデータは、4次元のTorch型でtorch.Size([1, 1, 28, 28])となっています。
(2) imshow()
numpyに変換してmatplotlibで表示すると以下のように、黒字に白で数字が書かれたデータであることが分かります。
(3) dump
ダンプしてみると、黒が0で白が1となる範囲で、表現されていることを確認できます。
tensor([[[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], ・・・省略・・・ [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3294, 0.7255, 0.6235, 0.5922, 0.2353, 0.1412, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], ・・・省略・・・ [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])
OpenCVで取得した画像を、上記の入力形式に合わせてやることで、モデルが利用可能になります。
4 使ってみる
モデルを使っているコードです。
OpenCVで取得した画像は、下記の変換を行って、モデルの入力としています。
- 推論の範囲を切り出す
- グレーススケールで白黒画像に変換する
- 白と黒を反転させる
- 2値化処理を行う
- サイズを28✕28に変換する
- 0〜255で表現されているデータを0〜1に変換する
- 次元を追加して、[1, 1, 28,28]に合わせる
import cv2 import time import torch import numpy as np import torch.nn as nn import torch.nn.functional as F # Webカメラ DEVICE_ID = 0 WIDTH = 800 HEIGHT = 600 path = './mnist_cnn.pt' class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) self.conv2 = nn.Conv2d(32, 64, 3, 1) self.dropout1 = nn.Dropout2d(0.25) self.dropout2 = nn.Dropout2d(0.5) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = F.relu(x) x = self.conv2(x) x = F.relu(x) x = F.max_pool2d(x, 2) x = self.dropout1(x) x = torch.flatten(x, 1) x = self.fc1(x) x = F.relu(x) x = self.dropout2(x) x = self.fc2(x) output = F.log_softmax(x, dim=1) return output def main(): cap = cv2.VideoCapture (DEVICE_ID) # フォーマット・解像度・FPSの設定 cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT) # フォーマット・解像度・FPSの取得 width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) fps = cap.get(cv2.CAP_PROP_FPS) print("fps:{} width:{} height:{}".format(fps, width, height)) # 検出範囲 size = 50 x1 = int(width/2-size) x2 = int(width/2+size) y1 = int(height/2-size) y2 = int(height/2+size) # モデル+パラメータの読み込み model = Net() model.load_state_dict(torch.load(path)) model.eval() # 評価モード while True: # カメラ画像取得 _, frame = cap.read() if(frame is None): continue img = frame[y1 : y2, x1 : x2] # 対象領域の抽出 img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # グレースケール img = cv2.bitwise_not(img) # 白黒反転 _, img = cv2.threshold(img, 0, 255, cv2.THRESH_OTSU) # 2値化 img = cv2.resize(img,(28, 28)) # サイズ変更 => 28 * 28 cv2.imshow('img', img) # モニター img = img/256 # 0〜255 => 0.0〜1.0 img = img[np.newaxis, np.newaxis, :, :] # 次元追加 (28,28) => (1, 1, 28, 28) pred = model(torch.tensor(img, dtype=torch.float32)) print(pred) index = pred.argmax() prediction = "Prediction: {}".format(pred.argmax().item()) print(prediction) cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 255, 0)) cv2.putText(frame, prediction, (10, int(height)-50), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 255, 255), 3, cv2.LINE_AA) # 画像表示 cv2.imshow('frame', frame) # 'q'をタイプされたらループから抜ける if cv2.waitKey(1) & 0xFF == ord('q'): break # VideoCaptureオブジェクト破棄 cap.release() cv2.destroyAllWindows() if __name__ == '__main__': main()
5 最後に
学習されたモデルを使用するために、入力形式を整合させることは必須の作業です。
作業としては、概ね、「縦、横、チャンネル」の順番や、ミニバッチの分の次元調整が主になると思います。 また、OpenCVでは、Numpy形式で画像を扱いますが、最終的にはTorch型にすることも忘れてはなりません。
次元やデータサイズが、モデルの入力に一致しないと、推論の段階でエラーとなるはずなので、そのエラーメッセージを見ることでも、だいだい間違いに気がつくかも知れません。